package cn.cosmosx.base.filter;

import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;
import cn.cosmosx.base.constant.ResultCode;
import cn.cosmosx.base.exception.PlatformException;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Map;

/**
 * XSS处理过滤器
 */
@Component
public class XSSFilter extends OncePerRequestFilter {
    // 请求类型：form表单提交的无需处理
    private final String[] UN_WRAPPER_LIST = {"multipart/form-data"};
    private final String[] XSS_STRS = {"<script", "<img", "<a href", "onerror=", "eval(", "=\"&#", "${jndi:"};

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain) throws ServletException, IOException {
        if (!(request instanceof HttpServletRequestWrapperFilter.RequestBodyWrapper)
                || (StringUtils.isNotEmpty(request.getContentType())
                && Arrays.stream(UN_WRAPPER_LIST).anyMatch(x -> request.getContentType().contains(x)))) {
            filterChain.doFilter(request, response);
            return;
        }
        checkParam((HttpServletRequestWrapperFilter.RequestBodyWrapper) request);
        filterChain.doFilter(request, response);
    }

    private void checkParam(HttpServletRequestWrapperFilter.RequestBodyWrapper request) {
        // check url param
        Map<String, String[]> map = request.getParameterMap();
        map.forEach((key, values) -> {
            for (String param : values) {
                matchXssStr(param);
            }
        });
        // check body param
        byte[] body = request.getBodyBuff();
        if (body == null || body.length == 0) {
            return;
        }
        matchXssStr(new String(body, StandardCharsets.UTF_8));
    }

    private void matchXssStr(String param) {
        if (StringUtils.isEmpty(param)) {
            return;
        }
        if (Arrays.stream(XSS_STRS).anyMatch(param::contains)) {
            throw new PlatformException(ResultCode.P00005);
        }
    }
}
